#!/usr/bin/env puython
import os
import sys
import re

## Usage: python pull_scores.py dir output
## If directory argument is not supplied, then current directory is used
## If output not supplied, then default is used

cur_dir = os.getcwd()
work_dir = cur_dir
output_file = 'pull_scores_output.txt'
PDB = '1a2k'
gz_search_string = '_seqtol.out'
gz_string_len = len(gz_search_string)

if len(sys.argv) > 1:
    work_dir = os.path.abspath(sys.argv[1])

if len(sys.argv) > 2:
    output_file = sys.argv[2]

score_list = []

for root, dirs, infiles in os.walk(work_dir):
    for infile in infiles:
        if infile[len(infile)-gz_string_len:]!=gz_search_string:
            continue
        print infile
        m = re.search('('+PDB+'[_])(\d*)('+gz_search_string+')',infile)
        best_fitness = float('Inf')
        best_backrub_score = float('Inf')
        num = 0
        if m:
            num = int(m.group(2))
            fileIn = open(os.path.join(root,infile),'r')
            linenum = 0
            for line in fileIn:
                #print best_fitness
                #print str(linenum) + ' ' + line
                linenum+=1
                m2 = re.search('(?:.*?)(?:with\sfitness[:]\s)(-?\d+[.]\d+)(?:.*)',line)
                if m2:
                    #print 'match'
                    fitness = float(m2.group(1))
                    if fitness < best_fitness:
                        best_fitness = fitness
            fileIn.close()
        if best_fitness < float('Inf'):
            fileIn = open(os.path.join(root,PDB+'_'+str(num)+'_backrub.out'))
            for line in fileIn:
                m3 =  re.search('(?:.*?)(?:Total\sweighted\sscore[:])'+
                                '(?:\s+?)(-?\d+?[.]\d+?)(?:.*?)',line)
                if m3:
                    backrub_score = float(m3.group(1))
                    if backrub_score < best_backrub_score:
                        best_backrub_score = backrub_score
            print str(num) + ' ' + str(best_backrub_score) + ' ' + str(best_fitness)
            new_tuple = num,best_backrub_score,best_fitness
            score_list.append(new_tuple)
print ''
print str(len(score_list)) + ' runs found. Now writing...'
fileOut = open(os.path.join(cur_dir,output_file),'w')
fileOut.write('RunNumber\tBestBackrubScore\tBestSeqTolScore\n')
for tup in score_list:
    fileOut.write(str(tup[0])+'\t'+str(tup[1])+'\t'+str(tup[2])+'\n')
print 'Done'
